package com.example.utils;
import com.example.dto.LoadRunnerResult;
import lombok.extern.slf4j.Slf4j;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/*
@description:压测工具类
@ClassName LoadRunnerUtils
@author chen
@create 2025-06-10 16:04
@Version 1.0
*/
@Slf4j
public class LoadRunnerUtils
{

    /*
     * @param  对 command 执行压测
     * @param requests 总请求数
     * @param concurrency 并发数量
     * @param command 需要执行的压测代码
     * @==============History===============<br/>;
     * @Description // 压测方法
     * @Date 2025/6/10
     */
    public static <T> LoadRunnerResult run(int requests, int concurrency, Runnable command) throws InterruptedException
    {
        log.info("压测开始......");
        ThreadPoolExecutor poolExecutor = new ThreadPoolExecutor(
                                                    concurrency,
                                                    concurrency,
                                                    0L,
                                                    TimeUnit.MILLISECONDS,
                                                    new LinkedBlockingQueue<Runnable>());
        poolExecutor.prestartAllCoreThreads();

        // 创建一个 CountDownLatch，用于阻塞当前线程池待所有请求处理完毕后，让当前线程继续向下走
        CountDownLatch countDownLatch = new CountDownLatch(requests);


        // 成功请求数、最快耗时、最慢耗时 (这几个值涉及到并发操作，所以采用 AtomicInteger 避免并发修改导致数据错误)
        AtomicInteger successRequests = new AtomicInteger(0);
        AtomicInteger fastestCostTime = new AtomicInteger(Integer.MAX_VALUE);
        AtomicInteger slowestCostTime = new AtomicInteger(Integer.MIN_VALUE);


        long startTime = System.currentTimeMillis();
        // 循环中使用线程池处理被压测的方法
        for (int i = 0; i < requests; i++)
        {
            poolExecutor.execute(()->
            {
                try
                {
                    long requestStartTime = System.currentTimeMillis();

                    command.run();

                    int costTime = (int) (System.currentTimeMillis() - requestStartTime);

                    // 请求请求最快耗时
                    setFastestCostTime(fastestCostTime, costTime);

                    // 请求最慢耗时
                    setSlowestCostTimeCostTime(slowestCostTime,costTime);

                    successRequests.incrementAndGet();
                }
                finally
                {
                    countDownLatch.countDown();
                }
            });
        }
        countDownLatch.await();
        poolExecutor.shutdown();

        long endTime = System.currentTimeMillis();
        log.info("压测结束，总耗时(ms):{}", (endTime - startTime));
        LoadRunnerResult result = new LoadRunnerResult();
        result.setRequests(requests);
        result.setConcurrency(concurrency);
        result.setSuccessRequests(successRequests.get());
        result.setFailRequests(requests - result.getSuccessRequests());
        result.setTimeTakenForTests((int) (endTime - startTime));
        result.setRequestsPerSecond((float) requests * 1000f / (float) (result.getTimeTakenForTests()));
        result.setTimePerRequest((float) result.getTimeTakenForTests() / (float) requests);
        result.setFastestCostTime(fastestCostTime.get());
        result.setSlowestCostTime(slowestCostTime.get());
        return result;
    }




    private static void setFastestCostTime(AtomicInteger fastestCostTime, int costTime)
    {
        while (true)
        {
            int fsCostTime = fastestCostTime.get();
            if (fsCostTime < costTime)
            {
                break;
            }
            if (fastestCostTime.compareAndSet(fsCostTime, costTime))
            {
                break;
            }
        }
    }


    private static void setSlowestCostTimeCostTime(AtomicInteger slowestCostTime, int costTime)
    {
        while (true) {
            int slCostTime = slowestCostTime.get();
            if (slCostTime > costTime)
            {
                break;
            }
            if (slowestCostTime.compareAndSet(slCostTime, costTime))
            {
                break;
            }
        }
    }

}
